{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "from scipy.integrate import odeint\n",
    "from utils.util import EMax,load_yaml_config,searchKey\n",
    "import pandas as pd\n",
    "import numpy as np\n",
    "\n",
    "class Compartment():\n",
    "    def __init__(self,name,parameters,clear_rate,out_rate,factors,*args,**kwargs):\n",
    "        self.name = name\n",
    "        self.factors = factors\n",
    "        self.parameters = parameters\n",
    "        self.clear_rate = clear_rate\n",
    "        self.out_rate = out_rate\n",
    "\n",
    "    def __getattr__(self, item):\n",
    "        if item in self.parameters.keys():\n",
    "            return self.parameters[item]\n",
    "        if item in self.factors:\n",
    "            return f'{self.name}.{item}'\n",
    "\n",
    "\n",
    "    def clear(self,y0):\n",
    "        clear = defaultdict(float)\n",
    "        if self.clear_rate is not None:\n",
    "            for key,value in self.clear_rate.items():\n",
    "                clear[key] =  value* self.y0[f'{self.name}.{key}']\n",
    "        return clear\n",
    "\n",
    "    @property\n",
    "    def delta(self,y0):\n",
    "        '''内部变化\n",
    "        '''\n",
    "        # 初始化梯度\n",
    "        delta = dict(zip([f'{self.name}.{f}' for f in self.factors],[0 for _ in self.factors]))\n",
    "        delta[self.drug] = self.kofftaa * y0[self.dimer_taa] + self.koffeaa * y0[self.dimer_eaa]  - y0[self.drug] * (self.kon_eaa * y0[self.eaa] + self.kontaa * y0[self.taa]) /self.volumn\n",
    "        delta[self.taa] = self.kofftaa * (y0[self.dimer_taa] + y0[self.trimer]) - kontaa * y0[self.taa] *(y0[self.drug] + y0[self.dimer_eaa])/self.volumn\n",
    "        delta[self.eaa] = self.koffeaa * (y0[self.dimer_eaa] + y0[self.trimer]) - koneaa * y0[self.eaa] *(y0[self.drug] + y0[self.dimer_taa])/self.volumn\n",
    "        delta[self.dimer_taa] = self.kontaa * y0[self.taa] * y0[self.drug]/self.volumn + self.koffeaa*y0[self.trimer] - self.kofftaa * y0[self.dimer_taa] - self.oneaa* y0[self.dimer_taa]*y0[self.eaa]\n",
    "        delta[self.dimer_eaa] =  self.koneaa * y0[self.eaa] * y0[self.drug]/self.volumn + self.kofftaa*y0[self.trimer] - self.koffeaa * y0[self.dimer_eaa] - self.ontaa* y0[self.dimer_eaa]*y0[self.taa]\n",
    "        delta[self.trimer] = self.oneaa* y0[self.dimer_taa]*y0[self.eaa] + self.ontaa* y0[self.dimer_eaa]*y0[self.taa] - self.kofftaa * y0[self.taa] - self.kofftaa*y0[self.trimer]\n",
    "        # Clear\n",
    "        if self.clear_rate is not None:\n",
    "            for key,value in self.clear_rate.items():\n",
    "                delta[key]  -=  value * self.y0[eval(f'self.{key}')]\n",
    "        # exchange\n",
    "        if self.out_rate is not None:\n",
    "            for key in self.out_rate:\n",
    "                for factor in self.out_rate[key].keys():\n",
    "                    rate = self.out_rate[key][factor] / self.volumn\n",
    "                    newvalue = y0[eval(f'self.{factor}')] * rate\n",
    "                    delta[f'{key}.{factor}'] += newvalue\n",
    "                    delta[eval(f'self.{factor}')] -= newvalue\n",
    "        return delta\n",
    "\n",
    "\n",
    "\n",
    "class QSPModel:\n",
    "    def __init__(self,config):\n",
    "        if type(config) is not dict:\n",
    "            self.config = load_yaml_config(config)\n",
    "        else:\n",
    "            self.config = config\n",
    "        self.output = None\n",
    "        self.compartments = []\n",
    "        self.columns = []\n",
    "        for com in self.compartments:\n",
    "            self.columns += [f'{com.name}.{f}' for f in com.factors]\n",
    "        self.columns_idx = {}\n",
    "        for i,col in enumerate(columns):\n",
    "            self.columns_idx[col] = i\n",
    "\n",
    "    def init_compartments(self)->None:\n",
    "        self.Compartments = dict()\n",
    "        for k,v in searchKey('Compartments',self.config).items():\n",
    "            self.Compartments[k] = Compartment(**v)\n",
    "        if searchKey('Emax',self.config):\n",
    "            self.Efficacy  = EMax()\n",
    "            self.Efficacy.load( searchKey('Emax',self.config))\n",
    "\n",
    "    def delta(self,y0,t,args):\n",
    "        delta = np.zeros(len(self.columns))\n",
    "        y0_dict = self.vector_todict(y0)\n",
    "        for com in self.compartments:\n",
    "            delta += self.dict_tovector(com.delta(y0_dict))\n",
    "\n",
    "        return delta\n",
    "\n",
    "    def simulation(self,y,t,*args,**kwargs):\n",
    "        ts = np.arange(self.interval,self.ed,self.interval)\n",
    "        self.output = odeint(self.simulation, y0, ts, args=())\n",
    "\n",
    "    def to_pandas(self):\n",
    "        return pd.DataFrame(self.output,self.columns)\n",
    "\n",
    "    def vector_todict(self,y):\n",
    "        assert len(y) ==len(self.columns)\n",
    "        return dict(zip(self.columns,y))\n",
    "\n",
    "    def dict_tovector(self,dict_y):\n",
    "        assert len(dict_y.keys())\n",
    "        return np.array([dict_y[k] for k in self.columns])\n",
    "\n",
    "\n",
    "\n"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
